from celery import Celery
from celery_app import celery_app
import time
import json
import requests
import os
from dotenv import load_dotenv
from typing import Dict, Any

# 加载环境变量
load_dotenv()

# 获取后端API地址
backend_url = os.getenv("BACKEND_URL", "http://localhost:8000")

@celery_app.task
def test_task():
    """测试任务"""
    return {"message": "Test task executed successfully"}

@celery_app.task(name="tasks.migrate_schema_task")
def migrate_schema_task(task_id: int, source_config: Dict[str, Any], target_config: Dict[str, Any]):
    """结构迁移任务"""
    try:
        print(f"开始执行任务 {task_id} 的结构迁移")
        print(f"源数据库: {source_config['database']} -> 目标数据库: {target_config['database']}")
        
        # 1. 连接源数据库
        print(f"连接源数据库: {source_config['type']}://{source_config['host']}:{source_config['port']}")
        source_conn = None
        if source_config['type'] == 'mysql':
            import pymysql
            source_conn = pymysql.connect(
                host=source_config['host'],
                port=source_config['port'],
                user=source_config['user'],
                password=source_config['password'],
                database=source_config['database'],
                charset='utf8mb4'
            )
        elif source_config['type'] == 'postgresql':
            import psycopg2
            source_conn = psycopg2.connect(
                host=source_config['host'],
                port=source_config['port'],
                user=source_config['user'],
                password=source_config['password'],
                database=source_config['database']
            )
        
        # 2. 连接目标数据库（如果不存在则创建）
        print(f"连接目标数据库: {target_config['type']}://{target_config['host']}:{target_config['port']}")
        target_conn = None
        try:
            if target_config['type'] == 'mysql':
                import pymysql
                target_conn = pymysql.connect(
                    host=target_config['host'],
                    port=target_config['port'],
                    user=target_config['user'],
                    password=target_config['password'],
                    database=target_config['database'],
                    charset='utf8mb4'
                )
            elif target_config['type'] == 'postgresql':
                import psycopg2
                target_conn = psycopg2.connect(
                    host=target_config['host'],
                    port=target_config['port'],
                    user=target_config['user'],
                    password=target_config['password'],
                    database=target_config['database']
                )
        except Exception as e:
            print(f"连接目标数据库失败: {e}")
            # 如果数据库不存在，尝试创建它
            if "Unknown database" in str(e) or "does not exist" in str(e):
                print(f"目标数据库 {target_config['database']} 不存在，正在创建...")
                if target_config['type'] == 'mysql':
                    import pymysql
                    # 先连接到默认数据库（通常是mysql）
                    temp_conn = pymysql.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        charset='utf8mb4'
                    )
                    with temp_conn.cursor() as cursor:
                        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {target_config['database']}")
                    temp_conn.close()
                    print(f"数据库 {target_config['database']} 创建成功")
                    
                    # 重新连接到新创建的数据库
                    target_conn = pymysql.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database=target_config['database'],
                        charset='utf8mb4'
                    )
                elif target_config['type'] == 'postgresql':
                    import psycopg2
                    # PostgreSQL需要通过postgres数据库连接来创建新数据库
                    temp_conn = psycopg2.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database='postgres'  # 连接到默认的postgres数据库
                    )
                    temp_conn.autocommit = True  # 创建数据库需要autocommit模式
                    with temp_conn.cursor() as cursor:
                        cursor.execute(f"CREATE DATABASE {target_config['database']}")
                    temp_conn.close()
                    print(f"数据库 {target_config['database']} 创建成功")
                    
                    # 重新连接到新创建的数据库
                    target_conn = psycopg2.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database=target_config['database']
                    )
            else:
                # 其他错误直接抛出
                raise e
        
        # 3. 获取源数据库结构
        print("获取源数据库结构信息")
        tables = []
        table_schemas = {}
        
        if source_config['type'] == 'mysql':
            with source_conn.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [row[0] for row in cursor.fetchall()]
                
                # 获取每个表的结构
                for table in tables:
                    cursor.execute(f"SHOW CREATE TABLE {table}")
                    create_table_sql = cursor.fetchone()[1]
                    table_schemas[table] = create_table_sql
        elif source_config['type'] == 'postgresql':
            with source_conn.cursor() as cursor:
                cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
                tables = [row[0] for row in cursor.fetchall()]
                
                # 获取每个表的结构
                for table in tables:
                    cursor.execute(f"SELECT pg_get_serial_sequence('{table}', column_name) FROM information_schema.columns WHERE table_name = '{table}'")
                    # 简化处理，实际项目中需要更复杂的逻辑
                    cursor.execute(f"SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_name = '{table}'")
                    columns = cursor.fetchall()
                    columns_def = ", ".join([f"{col[0]} {col[1].upper()}{' NULL' if col[2] == 'YES' else ' NOT NULL'}" for col in columns])
                    create_table_sql = f"CREATE TABLE {table} ({columns_def})"
                    table_schemas[table] = create_table_sql
        
        print(f"找到表: {tables}")
        
        # 4. 转换结构（根据目标数据库类型）
        print(f"将结构从 {source_config['type']} 转换为 {target_config['type']}")
        converted_schemas = {}
        for table, create_sql in table_schemas.items():
            # 简化的结构转换，实际项目中需要更复杂的逻辑
            if source_config['type'] == 'mysql' and target_config['type'] == 'postgresql':
                # 简单示例：替换部分关键字
                create_sql = create_sql.replace('`', '').replace('AUTO_INCREMENT', 'SERIAL')
            elif source_config['type'] == 'postgresql' and target_config['type'] == 'mysql':
                # 简单示例：替换部分关键字
                create_sql = create_sql.replace('SERIAL', 'INT AUTO_INCREMENT')
            
            converted_schemas[table] = create_sql
        
        # 5. 在目标数据库创建结构
        print("在目标数据库创建表结构")
        # 禁用外键约束检查（MySQL）
        if target_config['type'] == 'mysql':
            with target_conn.cursor() as cursor:
                cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
        
        try:
            for table, create_sql in converted_schemas.items():
                print(f"创建表 {table}")
                try:
                    with target_conn.cursor() as cursor:
                        # 先尝试删除表（如果存在）
                        try:
                            cursor.execute(f"DROP TABLE IF EXISTS {table}")
                        except Exception as drop_e:
                            print(f"删除表 {table} 失败: {drop_e}")
                        
                        # 创建新表
                        cursor.execute(create_sql)
                    target_conn.commit()
                    print(f"表 {table} 创建成功")
                except Exception as e:
                    print(f"创建表 {table} 失败: {e}")
                    # 继续处理其他表
        finally:
            # 重新启用外键约束检查（MySQL）
            if target_config['type'] == 'mysql':
                with target_conn.cursor() as cursor:
                    cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
        
        # 关闭数据库连接
        if source_conn:
            source_conn.close()
        if target_conn:
            target_conn.close()
        
        print(f"任务 {task_id} 的结构迁移完成")
        result = {
            "task_id": task_id,
            "status": "completed",
            "message": f"结构迁移完成，共迁移 {len(tables)} 个表"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
    except Exception as e:
        result = {
            "task_id": task_id,
            "status": "failed",
            "message": f"结构迁移失败: {str(e)}"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result

@celery_app.task(name="tasks.migrate_data_task")
def migrate_data_task(task_id: int, source_config: Dict[str, Any], target_config: Dict[str, Any], filter_condition: str = None):
    """数据迁移任务"""
    try:
        print(f"开始执行任务 {task_id} 的数据迁移")
        print(f"源数据库: {source_config['database']} -> 目标数据库: {target_config['database']}")
        
        # 1. 连接源数据库
        print(f"连接源数据库: {source_config['type']}://{source_config['host']}:{source_config['port']}")
        source_conn = None
        if source_config['type'] == 'mysql':
            import pymysql
            source_conn = pymysql.connect(
                host=source_config['host'],
                port=source_config['port'],
                user=source_config['user'],
                password=source_config['password'],
                database=source_config['database'],
                charset='utf8mb4'
            )
        elif source_config['type'] == 'postgresql':
            import psycopg2
            source_conn = psycopg2.connect(
                host=source_config['host'],
                port=source_config['port'],
                user=source_config['user'],
                password=source_config['password'],
                database=source_config['database']
            )
        
        # 2. 连接目标数据库（如果不存在则创建）
        print(f"连接目标数据库: {target_config['type']}://{target_config['host']}:{target_config['port']}")
        target_conn = None
        try:
            if target_config['type'] == 'mysql':
                import pymysql
                target_conn = pymysql.connect(
                    host=target_config['host'],
                    port=target_config['port'],
                    user=target_config['user'],
                    password=target_config['password'],
                    database=target_config['database'],
                    charset='utf8mb4'
                )
            elif target_config['type'] == 'postgresql':
                import psycopg2
                target_conn = psycopg2.connect(
                    host=target_config['host'],
                    port=target_config['port'],
                    user=target_config['user'],
                    password=target_config['password'],
                    database=target_config['database']
                )
        except Exception as e:
            print(f"连接目标数据库失败: {e}")
            # 如果数据库不存在，尝试创建它
            if "Unknown database" in str(e) or "does not exist" in str(e):
                print(f"目标数据库 {target_config['database']} 不存在，正在创建...")
                if target_config['type'] == 'mysql':
                    import pymysql
                    # 先连接到默认数据库（通常是mysql）
                    temp_conn = pymysql.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        charset='utf8mb4'
                    )
                    with temp_conn.cursor() as cursor:
                        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {target_config['database']}")
                    temp_conn.close()
                    print(f"数据库 {target_config['database']} 创建成功")
                    
                    # 重新连接到新创建的数据库
                    target_conn = pymysql.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database=target_config['database'],
                        charset='utf8mb4'
                    )
                elif target_config['type'] == 'postgresql':
                    import psycopg2
                    # PostgreSQL需要通过postgres数据库连接来创建新数据库
                    temp_conn = psycopg2.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database='postgres'  # 连接到默认的postgres数据库
                    )
                    temp_conn.autocommit = True  # 创建数据库需要autocommit模式
                    with temp_conn.cursor() as cursor:
                        cursor.execute(f"CREATE DATABASE {target_config['database']}")
                    temp_conn.close()
                    print(f"数据库 {target_config['database']} 创建成功")
                    
                    # 重新连接到新创建的数据库
                    target_conn = psycopg2.connect(
                        host=target_config['host'],
                        port=target_config['port'],
                        user=target_config['user'],
                        password=target_config['password'],
                        database=target_config['database']
                    )
            else:
                # 其他错误直接抛出
                raise e
        
        # 3. 获取表列表和依赖关系
        print("获取表列表和依赖关系")
        tables = []
        table_dependencies = {}  # 表依赖关系
        
        if source_config['type'] == 'mysql':
            with source_conn.cursor() as cursor:
                # 获取所有表
                cursor.execute("SHOW TABLES")
                tables = [row[0] for row in cursor.fetchall()]
                
                # 获取外键约束信息
                cursor.execute("""
                    SELECT 
                        TABLE_NAME,
                        REFERENCED_TABLE_NAME
                    FROM 
                        information_schema.KEY_COLUMN_USAGE 
                    WHERE 
                        REFERENCED_TABLE_NAME IS NOT NULL 
                        AND TABLE_SCHEMA = %s
                """, (source_config['database'],))
                
                for row in cursor.fetchall():
                    table_name, referenced_table = row
                    if table_name not in table_dependencies:
                        table_dependencies[table_name] = []
                    table_dependencies[table_name].append(referenced_table)
        elif source_config['type'] == 'postgresql':
            with source_conn.cursor() as cursor:
                # 获取所有表
                cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
                tables = [row[0] for row in cursor.fetchall()]
                
                # 获取外键约束信息
                cursor.execute("""
                    SELECT 
                        conrelid::regclass::text AS table_name,
                        confrelid::regclass::text AS referenced_table_name
                    FROM 
                        pg_constraint 
                    WHERE 
                        contype = 'f' 
                        AND connamespace = 'public'::regnamespace
                """)
                
                for row in cursor.fetchall():
                    table_name, referenced_table = row
                    if table_name not in table_dependencies:
                        table_dependencies[table_name] = []
                    table_dependencies[table_name].append(referenced_table)
        
        print(f"找到表: {tables}")
        print(f"表依赖关系: {table_dependencies}")
        
        # 4. 按依赖关系排序表
        sorted_tables = sort_tables_by_dependencies(tables, table_dependencies)
        print(f"排序后的表: {sorted_tables}")
        
        # 5. 逐表迁移数据（按依赖关系顺序）
        for table in sorted_tables:
            print(f"迁移表 {table} 的数据")
            
            # 获取源表数据
            with source_conn.cursor() as cursor:
                cursor.execute(f"SELECT * FROM {table}")
                columns = [desc[0] for desc in cursor.description]
                rows = cursor.fetchall()
                print(f"表 {table} 包含 {len(rows)} 行数据")
            
            # 在目标表中插入数据
            if rows:
                # 临时禁用外键约束（MySQL）
                if target_config['type'] == 'mysql':
                    with target_conn.cursor() as cursor:
                        cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
                
                try:
                    # 清空目标表数据
                    with target_conn.cursor() as cursor:
                        cursor.execute(f"DELETE FROM {table}")
                    
                    # 插入数据
                    if target_config['type'] == 'mysql':
                        if rows:
                            placeholders = ', '.join(['%s'] * len(columns))
                            insert_sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"
                            with target_conn.cursor() as cursor:
                                cursor.executemany(insert_sql, rows)
                    elif target_config['type'] == 'postgresql':
                        if rows:
                            placeholders = ', '.join([f"${i+1}" for i in range(len(columns))])
                            insert_sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"
                            with target_conn.cursor() as cursor:
                                cursor.executemany(insert_sql, rows)
                    
                    target_conn.commit()
                    print(f"表 {table} 数据迁移完成")
                except Exception as e:
                    # 发生错误时回滚
                    target_conn.rollback()
                    raise e
                finally:
                    # 重新启用外键约束（MySQL）
                    if target_config['type'] == 'mysql':
                        with target_conn.cursor() as cursor:
                            cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
            
        # 关闭数据库连接
        if source_conn:
            source_conn.close()
        if target_conn:
            target_conn.close()
        
        print(f"任务 {task_id} 的数据迁移完成")
        result = {
            "task_id": task_id,
            "status": "completed",
            "message": f"数据迁移完成，共迁移 {len(tables)} 个表"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
    except Exception as e:
        result = {
            "task_id": task_id,
            "status": "failed",
            "message": f"数据迁移失败: {str(e)}"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result

def sort_tables_by_dependencies(tables, dependencies):
    """
    根据表依赖关系对表进行排序，确保被依赖的表先迁移
    使用拓扑排序算法
    """
    # 初始化入度表
    in_degree = {table: 0 for table in tables}
    
    # 计算每个表的入度（被依赖次数）
    for table, deps in dependencies.items():
        for dep in deps:
            if dep in in_degree:
                in_degree[dep] += 1
    
    # 找到入度为0的表（没有被其他表依赖）
    queue = [table for table in tables if in_degree[table] == 0]
    result = []
    
    # 拓扑排序
    while queue:
        table = queue.pop(0)
        result.append(table)
        
        # 减少依赖该表的表的入度
        for dep_table, deps in dependencies.items():
            if table in deps:
                in_degree[dep_table] -= 1
                if in_degree[dep_table] == 0:
                    queue.append(dep_table)
    
    # 添加剩余的表（可能存在循环依赖）
    for table in tables:
        if table not in result:
            result.append(table)
    
    return result

@celery_app.task(name="tasks.sync_incremental_task")
def sync_incremental_task(task_id: int, source_config: Dict[str, Any], target_config: Dict[str, Any]):
    """增量同步任务"""
    try:
        print(f"开始执行任务 {task_id} 的增量同步")
        
        # 1. 连接源数据库
        print(f"连接源数据库: {source_config['type']}://{source_config['host']}:{source_config['port']}")
        time.sleep(1)
        
        # 2. 连接目标数据库
        print(f"连接目标数据库: {target_config['type']}://{target_config['host']}:{target_config['port']}")
        time.sleep(1)
        
        # 3. 持续监听变更并同步
        print("开始监听数据库变更...")
        for i in range(10):  # 模拟持续同步
            print(f"检查变更 #{i+1}")
            time.sleep(2)  # 模拟检查间隔
            
        print(f"任务 {task_id} 的增量同步完成")
        result = {
            "task_id": task_id,
            "status": "completed",
            "message": "增量同步完成"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
    except Exception as e:
        result = {
            "task_id": task_id,
            "status": "failed",
            "message": f"增量同步失败: {str(e)}"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result

@celery_app.task(name="tasks.migrate_mongodb_task")
def migrate_mongodb_task(task_id: int, source_config: Dict[str, Any], target_config: Dict[str, Any]):
    """MongoDB迁移任务"""
    try:
        print(f"开始执行任务 {task_id} 的MongoDB数据迁移")
        print(f"源数据库: {source_config['database']} -> 目标数据库: {target_config['database']}")
        
        # 1. 连接源MongoDB数据库
        print(f"连接源MongoDB数据库: mongodb://{source_config['host']}:{source_config['port']}")
        import pymongo
        source_client = pymongo.MongoClient(
            f"mongodb://{source_config['user']}:{source_config['password']}@{source_config['host']}:{source_config['port']}/"
        )
        source_db = source_client[source_config['database']]
        
        # 2. 连接目标MongoDB数据库
        print(f"连接目标MongoDB数据库: mongodb://{target_config['host']}:{target_config['port']}")
        target_client = pymongo.MongoClient(
            f"mongodb://{target_config['user']}:{target_config['password']}@{target_config['host']}:{target_config['port']}/"
        )
        target_db = target_client[target_config['database']]
        
        # 3. 获取所有集合（表）名称
        collections = source_db.list_collection_names()
        print(f"找到集合: {collections}")
        
        # 4. 逐个迁移集合
        for collection_name in collections:
            print(f"迁移集合 {collection_name}")
            collection = source_db[collection_name]
            
            # 删除目标集合中已有的数据
            target_db[collection_name].delete_many({})
            
            # 迁移数据
            documents = list(collection.find())
            if documents:
                target_db[collection_name].insert_many(documents)
            
            print(f"集合 {collection_name} 迁移完成，共迁移 {len(documents)} 条记录")
        
        # 关闭数据库连接
        source_client.close()
        target_client.close()
        
        print(f"任务 {task_id} 的MongoDB数据迁移完成")
        result = {
            "task_id": task_id,
            "status": "completed",
            "message": f"MongoDB数据迁移完成，共迁移 {len(collections)} 个集合"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
    except Exception as e:
        result = {
            "task_id": task_id,
            "status": "failed",
            "message": f"MongoDB数据迁移失败: {str(e)}"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result

@celery_app.task(name="tasks.migrate_redis_task")
def migrate_redis_task(task_id: int, source_config: Dict[str, Any], target_config: Dict[str, Any]):
    """Redis迁移任务"""
    try:
        print(f"开始执行任务 {task_id} 的Redis数据迁移")
        print(f"源数据库: {source_config['database']} -> 目标数据库: {target_config['database']}")
        
        # 1. 连接源Redis数据库
        print(f"连接源Redis数据库: {source_config['host']}:{source_config['port']}")
        import redis
        source_client = redis.Redis(
            host=source_config['host'],
            port=source_config['port'],
            password=source_config['password'] if source_config['password'] else None,
            db=int(source_config['database'].replace('db', '')) if source_config['database'].startswith('db') else 0
        )
        
        # 2. 连接目标Redis数据库
        print(f"连接目标Redis数据库: {target_config['host']}:{target_config['port']}")
        target_client = redis.Redis(
            host=target_config['host'],
            port=target_config['port'],
            password=target_config['password'] if target_config['password'] else None,
            db=int(target_config['database'].replace('db', '')) if target_config['database'].startswith('db') else 0
        )
        
        # 3. 迁移所有键值对
        keys = source_client.keys('*')
        print(f"找到 {len(keys)} 个键")
        
        # 逐个迁移键值对
        for key in keys:
            # 获取值和类型
            value = source_client.dump(key)
            ttl = source_client.ttl(key)
            
            # 在目标Redis中恢复键值对
            if value:
                target_client.restore(key, ttl if ttl > 0 else 0, value, replace=True)
        
        print(f"任务 {task_id} 的Redis数据迁移完成")
        result = {
            "task_id": task_id,
            "status": "completed",
            "message": f"Redis数据迁移完成，共迁移 {len(keys)} 个键"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
    except Exception as e:
        result = {
            "task_id": task_id,
            "status": "failed",
            "message": f"Redis数据迁移失败: {str(e)}"
        }
        
        # 通知后端更新任务状态
        try:
            requests.post(f"{backend_url}/api/v1/tasks/{task_id}/update-status", json=result)
        except Exception as e:
            print(f"通知后端更新状态失败: {e}")
            
        return result
